import os

import gensim.downloader as api
import numpy as np
import torch

path = os.path.join('.', 'data', 'sketch')
# part-shrec: 48
# label_list = [
#     'airplane', 'ant', 'armchair', 'bed', 'bench', 'bicycle', 'bookshelf', 'bus', 'cabinet', 'car', 'cell', 'chair',
#     'couch', 'door', 'eyeglasses', 'face', 'fish', 'floor', 'flower', 'flying', 'guitar', 'hand', 'head', 'helicopter',
#     'horse', 'house', 'keyboard', 'knife', 'laptop', 'motorbike', 'piano', 'potted', 'race', 'satellite', 'ship',
#     'spider', 'spoon', 'standing', 'sword', 'table', 'lamp', 'telephone', 'train', 'tree', 'truck', 'vase', 'violin',
#     'wheel'
# ]
# shrec14: 171
label_list = [
    'airplane', 'clock', 'ant', 'apple', 'armchair', 'ashtray', 'axe', 'banana', 'barn', 'baseball', 'basket',
    'bathtub', 'bear', 'bed', 'bee', 'mug', 'bench', 'bicycle', 'blimp', 'book', 'bookshelf', 'bowl', 'brain', 'bread',
    'bridge', 'bus', 'bush', 'butterfly', 'cabinet', 'cake', 'camel', 'camera', 'candle', 'cannon', 'car', 'castle',
    'phone', 'chair', 'chandelier', 'church', 'cigarette', 'monitor', 'couch', 'cow', 'crab', 'crocodile', 'cup',
    'diamond', 'dog', 'dolphin', 'door', 'handle', 'dragon', 'duck', 'ear', 'elephant', 'eyeglasses', 'face', 'fan',
    'hydrant', 'fish', 'lamp', 'flower', 'flying', 'saucer', 'fork', 'frog', 'giraffe', 'grapes', 'guitar', 'hammer',
    'hand', 'hat', 'head', 'helicopter', 'helmet', 'horse', 'balloon', 'hourglass', 'house', 'skeleton', 'ice',
    'kangaroo', 'key', 'keyboard', 'knife', 'ladder', 'laptop', 'leaf', 'lightbulb', 'lighter', 'lion', 'mailbox',
    'microphone', 'microscope', 'monkey', 'motorbike', 'mug', 'mushroom', 'octopus', 'owl', 'tree', 'parking', 'parrot',
    'pear', 'pen', 'penguin', 'sitting', 'walking', 'piano', 'truck', 'pig', 'smoking', 'plant', 'pumpkin', 'rabbit',
    'race', 'rifle', 'sailboat', 'satellite', 'dish', 'scissors', 'scorpion', 'screwdriver', 'turtle', 'seagull',
    'shark', 'sheep', 'ship', 'shoe', 'shovel', 'skateboard', 'skull', 'skyscraper', 'snake', 'snowman', 'shuttle',
    'spider', 'spoon', 'bird', 'stapler', 'submarine', 'suitcase', 'suv', 'sword', 'syringe', 'table', 'tablelamp',
    'teacup', 'teapot', 'teddy', 'telephone', 'tennis', 'tent', 'tire', 'toilet', 'tooth', 'light', 'train', 'tree',
    'truck', 'tv', 'umbrella', 'van', 'vase', 'violin', 'wheel', 'windmill', 'bottle', 'wineglass', 'watch'
]
print(label_list)
w2v = []

google_300_corpus = api.load("word2vec-google-news-300")
for label in label_list:
    print(label, 'done')
    w2v.append(torch.from_numpy(np.array(google_300_corpus[label])))

print(w2v, len(w2v))
torch.save(w2v, os.path.join('.', 'pth', 'label_word_vector.pth'))
